{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from charm.toolbox.pairinggroup import PairingGroup, G1, ZR, pair as e\n",
    "from Util import read_word_as_bytes\n",
    "from functools import reduce\n",
    "from hashlib import sha256, md5\n",
    "\n",
    "group = PairingGroup('SS512')\n",
    "g = group.random(G1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$ H_1: \\rightarrow G_1 $\\\n",
    "$ H_2: \\rightarrow G_1 $ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def H1(b):\n",
    "    return group.hash(sha256(b).hexdigest(), G1)\n",
    "\n",
    "def H2(b):\n",
    "    return group.hash(md5(b).hexdigest(), G1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "setup \\\n",
    "数据检索则选取随机数$x = Z_p^*$以及群$G_1$的生成元$g$, 最后得到公钥 $ pk = y $ 以及私钥 $ sk = x $ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup():\n",
    "    x = group.random(ZR)\n",
    "    y = g ** x\n",
    "    return {'pk': y, 'sk': x}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "build_index \\\n",
    "假设关键词集合 $ W = \\{w_1, w_2, ..., w_l\\} $ 生成索引，数据所有者首先选取两个随机数，$ s, r \\in Z_p^* $，然后根据公钥和随机数计算 $ A = g^r, B = y^s, C_i = h_i^rf_i^s(1 \\leq i \\leq l) $ ，其中 $h_i = H_1(w_i), f_i = H_2(w_i)$。最后，输出关键词集合对应的索引 $ I = \\{A, B, C_1, C_2, ..., C_l\\} $  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_index(W, pk):\n",
    "    s, r = group.random(ZR, 2)\n",
    "    return g ** r, pk ** s, [(H1(w_i) ** r) * (H2(w_i) ** s) for w_i in W]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "generate_trapdoor \\\n",
    "对于查询条件$ W=\\left\\{ ind_1,ind_2,...,ind_n,w_{ind_1},w_{ind_2},...,w_{ind_n} \\right\\}  $，数据检索者选取随机数 $ t \\in Z_p^* $ ，并计算$ T_1=g^t,T_2=\\left( h_{ind_i},h_{ind_{i+1}},...,h_{ind_m} \\right) ^t,T_3=\\left( f_{ind_i},h_{ind_{i+1}},...,h_{ind_m} \\right) ^{t/x} $。随后输出陷门$ T=\\left\\{ T_1,T_2,T_3,ind_1,ind_2,...,ind_m \\right\\}  $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_trapdoor(Q, sk):\n",
    "    L = len(Q)\n",
    "    if L % 2 != 0:\n",
    "        raise ValueError('the length of query list must be an even number')\n",
    "    Qi, Qw = Q[:L // 2], Q[L // 2:]\n",
    "    t = group.random(ZR)\n",
    "    T1 = g ** t\n",
    "    T2 = reduce(lambda h_i, h_i_1: h_i * h_i_1, [H1(wi) for wi in Qw]) ** t\n",
    "    T3 = reduce(lambda f_i, f_i_1: f_i * f_i_1, [H2(wi) for wi in Qw]) ** (t / sk)\n",
    "    return T1, T2, T3, Qi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "search \\\n",
    "给定索引 $ I $ 和陷门 $ T $ ，测试是否满足\n",
    "$$\n",
    "e\\left( T_1,\\prod_{i=1}^m{C_{ind_i}} \\right) =e\\left( A,T_2 \\right) \\cdot e\\left( B,T_3 \\right)\n",
    "$$\n",
    "正确性证明：\n",
    "\\begin{align*}\n",
    "\\text{右边}&=e\\left( g^r,\\left( h_{ind_i}h_{ind_{i+1}}\\cdots h_{ind_m} \\right) ^t \\right) \\cdot e\\left( g^{xs},\\left( f_{ind_i}f_{ind_{i+1}}\\cdots f_{ind_m} \\right) ^{t/x} \\right) \n",
    "\\\\\n",
    "&=e\\left( g^t,\\left( h_{ind_i}h_{ind_{i+1}}\\cdots h_{ind_m} \\right) ^r \\right) \\cdot e\\left( g^t,\\left( f_{ind_i}f_{ind_{i+1}}\\cdots f_{ind_m} \\right) ^{\\left( xs \\right) /x} \\right) \n",
    "\\\\\n",
    "&=e\\left( g^t,\\left( h_{ind_i}h_{ind_{i+1}}\\cdots h_{ind_m} \\right) ^r\\cdot \\left( f_{ind_i}f_{ind_{i+1}}\\cdots f_{ind_m} \\right) ^s \\right) \n",
    "\\\\\n",
    "&=e\\left( g^t,\\left( {h_{ind_i}}^r{f_{ind_i}}^s \\right) \\left( {h_{ind_{i+1}}}^r{f_{ind_{i+1}}}^s \\right) \\cdots \\left( {h_{ind_m}}^r{f_{ind_m}}^s \\right) \\right) \n",
    "\\\\\n",
    "&=e\\left( g^t,\\prod_{i=1}^m{C_{ind_i}} \\right) \n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def search(Idx, T):\n",
    "    A, B, C = Idx\n",
    "    T1, T2, T3, Qi = T\n",
    "    return e(T1, reduce(lambda C_ind_i, C_ind_i_1: C_ind_i * C_ind_i_1, [C[ind_i] for ind_i in Qi])) == e(A, T2) * e(\n",
    "        B, T3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "key_pair:  {'pk': [7159964578448490221234726169553013180803958811836429885265639467462975611651246770986995289592289750384064515548294581639229376242056021250334109562011286, 555674336510452353240648725399558118423664712577357329683103977131156183230968472468088271105708796438476460507504710834875486017105424254778196006214977], 'sk': 167613915869075949459448169478213201072178899044}\n",
      "search list: [0, 1, 5, b'There', b'is', b'happiness'], result: Yes\n"
     ]
    }
   ],
   "source": [
    "world_list = read_word_as_bytes('../test_files/happiness.txt')\n",
    "Q = [0, 1, 5, b'There', b'is', b'happiness']\n",
    "key_pair = setup()\n",
    "print('key_pair: ', key_pair)\n",
    "Idx = build_index(world_list, key_pair['pk'])\n",
    "T = generate_trapdoor(Q, key_pair['sk'])\n",
    "print(f'search list: {Q}, result: ', end = '')\n",
    "if search(Idx, T):\n",
    "    print('Yes')\n",
    "else:\n",
    "    print('No')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.9 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "d7eccc50bd8023cf74f021c24ed3740245f8543f295cddc81d7590622531cffd"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
